# 图片操作
from PIL import Image
import numpy as np
import cv2
from zz import zz_return


def find_img(img_src: Image.Image, img_target: Image.Image, threshold=0.9, get_index=0) -> zz_return.Return:
    """
    目标图在原图的位置,返回x,y
    :param img_src: 原图
    :param img_target: 目标图
    :param threshold: 设置阈值
    :param get_index: 匹配成功的数组下标，默认获取第1个
    :return:
    """
    # 图片灰度模式读取,为了速度更快
    img_arr = np.array(img_src.convert('L'), 'f')  # 以灰度模式读取
    img_find_arr = np.array(img_target.convert('L'), 'f')  # 以灰度模式读取

    # 获取模板的宽度和高度
    # w, h = img_find_arr.shape[::-1]

    # 归一化相关系数匹配法
    # res = cv2.matchTemplate(img_arr, img_find_arr, cv2.TM_CCOEFF_NORMED)  # 执行模板匹配
    # loc = np.where(res >= threshold)  # 查找最佳匹配位置

    # 使用平方差匹配算法
    res = cv2.matchTemplate(img_arr, img_find_arr, cv2.TM_SQDIFF)
    loc = np.where(res < threshold)
    loc = list(zip(*loc[::-1]))  # 将两个1维组数，合并变成（x,y）
    # print(loc)
    if len(loc) <= get_index:
        return zz_return.of(20029, "未匹配到位置")

    # 目标图第1个像素的位置
    first_x, first_y = loc[get_index]
    first_x = int(first_x)
    first_y = int(first_y)
    # 计算出点击中心位置
    center_x = first_x + int(img_target.width / 2)
    center_y = first_y + int(img_target.height / 2)

    r = zz_return.of(0)
    r.set_data("first", [first_x, first_y])
    r.set_data("center", [center_x, center_y])
    r.set_data("match_count", len(loc))

    # 返回匹配位置
    return r


def find_color(img_src, color, get_index=0) -> zz_return.Return:
    # img_arr = np.array(img_src.convert('L'))
    np_color = np.array(color)
    loc = np.where(img_src == np_color)
    if len(loc[0]) == 0:
        return zz_return.of(20053, "未找到匹配的位置")

    first_x = int(loc[1][get_index])
    first_y = int(loc[0][get_index])
    r = zz_return.of(0)
    r.set_data("location", [first_x, first_y])
    return r


def img_to_color_arr(img_src: Image.Image):
    """
    图片转颜色矩阵
    :param img_src:
    """
    color_arr = []
    width, height = img_src.size
    for y in range(height):
        color_item = []
        for x in range(width):
            pixel = img_src.getpixel((x, y))
            color_item.append(pixel)

        color_arr.append(color_item)

    return color_arr


def color_eq(color_a, color_b, tolerance=0):
    """
    两个颜色相同
    :param color_a: 格式(224, 224, 224)
    :param color_b: 格式(224, 224, 224)
    :param tolerance: 容忍度，允许颜色值有一定的偏差
    :return:
    """
    return all(abs(color_a[i] - color_b[i]) <= tolerance for i in range(3))
